/*
 * Copyright (c) 2021 Nordic Semiconductor ASA
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#define DT_DRV_COMPAT nordic_nrf_i2s

#include <stdlib.h>
#include <zephyr/drivers/i2s.h>
#include <zephyr/drivers/clock_control/nrf_clock_control.h>
#include <zephyr/drivers/pinctrl.h>
#include <soc.h>
#include <nrfx_i2s.h>
#include <hal/nrf_clock.h>

#include <zephyr/logging/log.h>
#include <zephyr/irq.h>
LOG_MODULE_REGISTER(i2s_nrfx, CONFIG_I2S_LOG_LEVEL);

struct stream_cfg {
	struct i2s_config cfg;
	nrfx_i2s_config_t nrfx_cfg;
};

struct i2s_buf {
	void *mem_block;
	size_t size;
};

struct i2s_nrfx_drv_data {
	struct onoff_manager *clk_mgr;
	struct onoff_client clk_cli;
	struct stream_cfg tx;
	struct k_msgq tx_queue;
	struct stream_cfg rx;
	struct k_msgq rx_queue;
	nrfx_i2s_t i2s;
	const uint32_t *last_tx_buffer;
	enum i2s_state state;
	enum i2s_dir active_dir;
	bool stop;       /* stop after the current (TX or RX) block */
	bool discard_rx; /* discard further RX blocks */
	volatile bool next_tx_buffer_needed;
	bool tx_configured: 1;
	bool rx_configured: 1;
	bool request_clock: 1;
};

struct i2s_nrfx_drv_cfg {
	nrfx_i2s_data_handler_t data_handler;
	nrfx_i2s_config_t nrfx_def_cfg;
	const struct pinctrl_dev_config *pcfg;
	enum clock_source {
		PCLK32M,
		PCLK32M_HFXO,
		ACLK
	} clk_src;
};

/* Finds the clock settings that give the frame clock frequency closest to
 * the one requested, taking into account the hardware limitations.
 */
static void find_suitable_clock(const struct i2s_nrfx_drv_cfg *drv_cfg,
				nrfx_i2s_config_t *config,
				const struct i2s_config *i2s_cfg)
{
	const nrfx_i2s_clk_params_t clk_params = {
		.base_clock_freq =
			(NRF_I2S_HAS_CLKCONFIG && drv_cfg->clk_src == ACLK)
			/* The I2S_NRFX_DEVICE() macro contains build assertions that
			 * make sure that the ACLK clock source is only used when it is
			 * available and only with the "hfclkaudio-frequency" property
			 * defined, but the default value of 0 here needs to be used to
			 * prevent compilation errors when the property is not defined
			 * (this expression will be eventually optimized away then).
			 */
			? DT_PROP_OR(DT_NODELABEL(clock), hfclkaudio_frequency, 0)
			: 32*1000*1000UL,
		.transfer_rate = i2s_cfg->frame_clk_freq,
		.swidth = config->sample_width,
		.allow_bypass = IS_ENABLED(CONFIG_I2S_NRFX_ALLOW_MCK_BYPASS),
	};

	if (nrfx_i2s_prescalers_calc(&clk_params, &config->prescalers) != 0) {
		LOG_ERR("Failed to find suitable I2S clock configuration.");
	}
}

static bool get_next_tx_buffer(struct i2s_nrfx_drv_data *drv_data,
			       nrfx_i2s_buffers_t *buffers)
{
	struct i2s_buf buf;
	int ret = k_msgq_get(&drv_data->tx_queue,
			     &buf,
			     K_NO_WAIT);
	if (ret == 0) {
		buffers->p_tx_buffer = buf.mem_block;
		buffers->buffer_size = buf.size / sizeof(uint32_t);
	}
	return (ret == 0);
}

static bool get_next_rx_buffer(struct i2s_nrfx_drv_data *drv_data,
			       nrfx_i2s_buffers_t *buffers)
{
	int ret = k_mem_slab_alloc(drv_data->rx.cfg.mem_slab,
				   (void **)&buffers->p_rx_buffer,
				   K_NO_WAIT);
	if (ret < 0) {
		LOG_ERR("Failed to allocate next RX buffer: %d",
			ret);
		return false;
	}

	return true;
}

static void free_tx_buffer(struct i2s_nrfx_drv_data *drv_data,
			   const void *buffer)
{
	k_mem_slab_free(drv_data->tx.cfg.mem_slab, (void *)buffer);
	LOG_DBG("Freed TX %p", buffer);
}

static void free_rx_buffer(struct i2s_nrfx_drv_data *drv_data, void *buffer)
{
	k_mem_slab_free(drv_data->rx.cfg.mem_slab, buffer);
	LOG_DBG("Freed RX %p", buffer);
}

static bool supply_next_buffers(struct i2s_nrfx_drv_data *drv_data,
				nrfx_i2s_buffers_t *next)
{
	if (drv_data->active_dir != I2S_DIR_TX) { /* -> RX active */
		if (!get_next_rx_buffer(drv_data, next)) {
			drv_data->state = I2S_STATE_ERROR;
			nrfx_i2s_stop(&drv_data->i2s);
			return false;
		}
		/* Set buffer size if there is no TX buffer (which effectively
		 * controls how many bytes will be received).
		 */
		if (drv_data->active_dir == I2S_DIR_RX) {
			next->buffer_size =
				drv_data->rx.cfg.block_size / sizeof(uint32_t);
		}
	}

	drv_data->last_tx_buffer = next->p_tx_buffer;

	LOG_DBG("Next buffers: %p/%p", next->p_tx_buffer, next->p_rx_buffer);
	nrfx_i2s_next_buffers_set(&drv_data->i2s, next);
	return true;
}

static void data_handler(const struct device *dev,
			 const nrfx_i2s_buffers_t *released, uint32_t status)
{
	struct i2s_nrfx_drv_data *drv_data = dev->data;
	bool stop_transfer = false;

	if (status & NRFX_I2S_STATUS_TRANSFER_STOPPED) {
		if (drv_data->state == I2S_STATE_STOPPING) {
			drv_data->state = I2S_STATE_READY;
		}
		if (drv_data->last_tx_buffer) {
			/* Usually, these pointers are equal, i.e. the last TX
			 * buffer that were to be transferred is released by the
			 * driver after it stops. The last TX buffer pointer is
			 * then set to NULL here so that the buffer can be freed
			 * below, just as any other TX buffer released by the
			 * driver. However, it may happen that the buffer is not
			 * released this way, for example, when the transfer
			 * ends with an error because an RX buffer allocation
			 * fails. In such case, the last TX buffer needs to be
			 * freed here.
			 */
			if (drv_data->last_tx_buffer != released->p_tx_buffer) {
				free_tx_buffer(drv_data,
					       drv_data->last_tx_buffer);
			}
			drv_data->last_tx_buffer = NULL;
		}
		nrfx_i2s_uninit(&drv_data->i2s);
		if (drv_data->request_clock) {
			(void)onoff_release(drv_data->clk_mgr);
		}
	}

	if (released == NULL) {
		/* This means that buffers for the next part of the transfer
		 * were not supplied and the previous ones cannot be released
		 * yet, as pointers to them were latched in the I2S registers.
		 * It is not an error when the transfer is to be stopped (those
		 * buffers will be released after the transfer actually stops).
		 */
		if (drv_data->state != I2S_STATE_STOPPING) {
			LOG_ERR("Next buffers not supplied on time");
			drv_data->state = I2S_STATE_ERROR;
		}
		nrfx_i2s_stop(&drv_data->i2s);
		return;
	}

	if (released->p_rx_buffer) {
		if (drv_data->discard_rx) {
			free_rx_buffer(drv_data, released->p_rx_buffer);
		} else {
			struct i2s_buf buf = {
				.mem_block = released->p_rx_buffer,
				.size = released->buffer_size * sizeof(uint32_t)
			};
			int ret = k_msgq_put(&drv_data->rx_queue,
					     &buf,
					     K_NO_WAIT);
			if (ret < 0) {
				LOG_ERR("No room in RX queue");
				drv_data->state = I2S_STATE_ERROR;
				stop_transfer = true;

				free_rx_buffer(drv_data, released->p_rx_buffer);
			} else {
				LOG_DBG("Queued RX %p", released->p_rx_buffer);

				/* If the TX direction is not active and
				 * the transfer should be stopped after
				 * the current block, stop the reception.
				 */
				if (drv_data->active_dir == I2S_DIR_RX &&
				    drv_data->stop) {
					drv_data->discard_rx = true;
					stop_transfer = true;
				}
			}
		}
	}

	if (released->p_tx_buffer) {
		/* If the last buffer that was to be transferred has just been
		 * released, it is time to stop the transfer.
		 */
		if (released->p_tx_buffer == drv_data->last_tx_buffer) {
			drv_data->discard_rx = true;
			stop_transfer = true;
		} else {
			free_tx_buffer(drv_data, released->p_tx_buffer);
		}
	}

	if (stop_transfer) {
		nrfx_i2s_stop(&drv_data->i2s);
	} else if (status & NRFX_I2S_STATUS_NEXT_BUFFERS_NEEDED) {
		nrfx_i2s_buffers_t next = { 0 };

		if (drv_data->active_dir != I2S_DIR_RX) { /* -> TX active */
			if (drv_data->stop) {
				/* If the stream is to be stopped, don't get
				 * the next TX buffer from the queue, instead
				 * supply the one used last time (it won't be
				 * transferred, the stream will stop right
				 * before this buffer would be started again).
				 */
				next.p_tx_buffer = drv_data->last_tx_buffer;
				next.buffer_size = 1;
			} else if (get_next_tx_buffer(drv_data, &next)) {
				/* Next TX buffer successfully retrieved from
				 * the queue, nothing more to do here.
				 */
			} else if (drv_data->state == I2S_STATE_STOPPING) {
				/* If there are no more TX blocks queued and
				 * the current state is STOPPING (so the DRAIN
				 * command was triggered) it is time to finish
				 * the transfer.
				 */
				drv_data->stop = true;
				/* Supply the same buffer as last time; it will
				 * not be transferred anyway, as the transfer
				 * will be stopped earlier.
				 */
				next.p_tx_buffer = drv_data->last_tx_buffer;
				next.buffer_size = 1;
			} else {
				/* Next TX buffer cannot be supplied now.
				 * Defer it to when the user writes more data.
				 */
				drv_data->next_tx_buffer_needed = true;
				return;
			}
		}

		(void)supply_next_buffers(drv_data, &next);
	}
}

static void purge_queue(const struct device *dev, enum i2s_dir dir)
{
	struct i2s_nrfx_drv_data *drv_data = dev->data;
	struct i2s_buf buf;

	if (dir == I2S_DIR_TX || dir == I2S_DIR_BOTH) {
		while (k_msgq_get(&drv_data->tx_queue,
				  &buf,
				  K_NO_WAIT) == 0) {
			free_tx_buffer(drv_data, buf.mem_block);
		}
	}

	if (dir == I2S_DIR_RX || dir == I2S_DIR_BOTH) {
		while (k_msgq_get(&drv_data->rx_queue,
				  &buf,
				  K_NO_WAIT) == 0) {
			free_rx_buffer(drv_data, buf.mem_block);
		}
	}
}

static int i2s_nrfx_configure(const struct device *dev, enum i2s_dir dir,
			      const struct i2s_config *i2s_cfg)
{
	struct i2s_nrfx_drv_data *drv_data = dev->data;
	const struct i2s_nrfx_drv_cfg *drv_cfg = dev->config;
	nrfx_i2s_config_t nrfx_cfg;

	if (drv_data->state != I2S_STATE_READY) {
		LOG_ERR("Cannot configure in state: %d", drv_data->state);
		return -EINVAL;
	}

	if (i2s_cfg->frame_clk_freq == 0) { /* -> reset state */
		purge_queue(dev, dir);
		if (dir == I2S_DIR_TX || dir == I2S_DIR_BOTH) {
			drv_data->tx_configured = false;
			memset(&drv_data->tx, 0, sizeof(drv_data->tx));
		}
		if (dir == I2S_DIR_RX || dir == I2S_DIR_BOTH) {
			drv_data->rx_configured = false;
			memset(&drv_data->rx, 0, sizeof(drv_data->rx));
		}
		return 0;
	}

	__ASSERT_NO_MSG(i2s_cfg->mem_slab != NULL &&
			i2s_cfg->block_size != 0);

	if ((i2s_cfg->block_size % sizeof(uint32_t)) != 0) {
		LOG_ERR("This device can transfer only full 32-bit words");
		return -EINVAL;
	}

	nrfx_cfg = drv_cfg->nrfx_def_cfg;

	switch (i2s_cfg->word_size) {
	case 8:
		nrfx_cfg.sample_width = NRF_I2S_SWIDTH_8BIT;
		break;
	case 16:
		nrfx_cfg.sample_width = NRF_I2S_SWIDTH_16BIT;
		break;
	case 24:
		nrfx_cfg.sample_width = NRF_I2S_SWIDTH_24BIT;
		break;
#if defined(I2S_CONFIG_SWIDTH_SWIDTH_32Bit)
	case 32:
		nrfx_cfg.sample_width = NRF_I2S_SWIDTH_32BIT;
		break;
#endif
	default:
		LOG_ERR("Unsupported word size: %u", i2s_cfg->word_size);
		return -EINVAL;
	}

	switch (i2s_cfg->format & I2S_FMT_DATA_FORMAT_MASK) {
	case I2S_FMT_DATA_FORMAT_I2S:
		nrfx_cfg.alignment = NRF_I2S_ALIGN_LEFT;
		nrfx_cfg.format = NRF_I2S_FORMAT_I2S;
		break;
	case I2S_FMT_DATA_FORMAT_LEFT_JUSTIFIED:
		nrfx_cfg.alignment = NRF_I2S_ALIGN_LEFT;
		nrfx_cfg.format = NRF_I2S_FORMAT_ALIGNED;
		break;
	case I2S_FMT_DATA_FORMAT_RIGHT_JUSTIFIED:
		nrfx_cfg.alignment = NRF_I2S_ALIGN_RIGHT;
		nrfx_cfg.format = NRF_I2S_FORMAT_ALIGNED;
		break;
	default:
		LOG_ERR("Unsupported data format: 0x%02x", i2s_cfg->format);
		return -EINVAL;
	}

	if ((i2s_cfg->format & I2S_FMT_DATA_ORDER_LSB) ||
	    (i2s_cfg->format & I2S_FMT_BIT_CLK_INV) ||
	    (i2s_cfg->format & I2S_FMT_FRAME_CLK_INV)) {
		LOG_ERR("Unsupported stream format: 0x%02x", i2s_cfg->format);
		return -EINVAL;
	}

	if (i2s_cfg->channels == 2) {
		nrfx_cfg.channels = NRF_I2S_CHANNELS_STEREO;
	} else if (i2s_cfg->channels == 1) {
		nrfx_cfg.channels = NRF_I2S_CHANNELS_LEFT;
	} else {
		LOG_ERR("Unsupported number of channels: %u",
			i2s_cfg->channels);
		return -EINVAL;
	}

	if ((i2s_cfg->options & I2S_OPT_BIT_CLK_SLAVE) &&
	    (i2s_cfg->options & I2S_OPT_FRAME_CLK_SLAVE)) {
		nrfx_cfg.mode = NRF_I2S_MODE_SLAVE;
	} else if (!(i2s_cfg->options & I2S_OPT_BIT_CLK_SLAVE) &&
		   !(i2s_cfg->options & I2S_OPT_FRAME_CLK_SLAVE)) {
		nrfx_cfg.mode = NRF_I2S_MODE_MASTER;
	} else {
		LOG_ERR("Unsupported operation mode: 0x%02x", i2s_cfg->options);
		return -EINVAL;
	}

	/* If the master clock generator is needed (i.e. in Master mode or when
	 * the MCK output is used), find a suitable clock configuration for it.
	 */
	if (nrfx_cfg.mode == NRF_I2S_MODE_MASTER ||
	    nrf_i2s_mck_pin_connected_check(drv_data->i2s.p_reg)) {
		find_suitable_clock(drv_cfg, &nrfx_cfg, i2s_cfg);
		/* Unless the PCLK32M source is used with the HFINT oscillator
		 * (which is always available without any additional actions),
		 * it is required to request the proper clock to be running
		 * before starting the transfer itself.
		 */
		drv_data->request_clock = (drv_cfg->clk_src != PCLK32M);
	} else {
		nrfx_cfg.prescalers.mck_setup = NRF_I2S_MCK_DISABLED;
		drv_data->request_clock = false;
	}

	if ((i2s_cfg->options & I2S_OPT_LOOPBACK) ||
	    (i2s_cfg->options & I2S_OPT_PINGPONG)) {
		LOG_ERR("Unsupported options: 0x%02x", i2s_cfg->options);
		return -EINVAL;
	}

	if (dir == I2S_DIR_TX || dir == I2S_DIR_BOTH) {
		drv_data->tx.cfg = *i2s_cfg;
		drv_data->tx.nrfx_cfg = nrfx_cfg;
		drv_data->tx_configured = true;
	}

	if (dir == I2S_DIR_RX || dir == I2S_DIR_BOTH) {
		drv_data->rx.cfg = *i2s_cfg;
		drv_data->rx.nrfx_cfg = nrfx_cfg;
		drv_data->rx_configured = true;
	}

	return 0;
}

static const struct i2s_config *i2s_nrfx_config_get(const struct device *dev,
						    enum i2s_dir dir)
{
	struct i2s_nrfx_drv_data *drv_data = dev->data;

	if (dir == I2S_DIR_TX && drv_data->tx_configured) {
		return &drv_data->tx.cfg;
	}
	if (dir == I2S_DIR_RX && drv_data->rx_configured) {
		return &drv_data->rx.cfg;
	}

	return NULL;
}

static int i2s_nrfx_read(const struct device *dev,
			 void **mem_block, size_t *size)
{
	struct i2s_nrfx_drv_data *drv_data = dev->data;
	struct i2s_buf buf;
	int ret;

	if (!drv_data->rx_configured) {
		LOG_ERR("Device is not configured");
		return -EIO;
	}

	ret = k_msgq_get(&drv_data->rx_queue,
			 &buf,
			 (drv_data->state == I2S_STATE_ERROR)
				? K_NO_WAIT
				: SYS_TIMEOUT_MS(drv_data->rx.cfg.timeout));
	if (ret == -ENOMSG) {
		return -EIO;
	}

	LOG_DBG("Released RX %p", buf.mem_block);

	if (ret == 0) {
		*mem_block = buf.mem_block;
		*size = buf.size;
	}

	return ret;
}

static int i2s_nrfx_write(const struct device *dev,
			  void *mem_block, size_t size)
{
	struct i2s_nrfx_drv_data *drv_data = dev->data;
	struct i2s_buf buf = { .mem_block = mem_block, .size = size };
	int ret;

	if (!drv_data->tx_configured) {
		LOG_ERR("Device is not configured");
		return -EIO;
	}

	if (drv_data->state != I2S_STATE_RUNNING &&
	    drv_data->state != I2S_STATE_READY) {
		LOG_ERR("Cannot write in state: %d", drv_data->state);
		return -EIO;
	}

	if (size > drv_data->tx.cfg.block_size || size < sizeof(uint32_t)) {
		LOG_ERR("This device can only write blocks up to %u bytes",
			drv_data->tx.cfg.block_size);
		return -EIO;
	}

	ret = k_msgq_put(&drv_data->tx_queue,
			 &buf,
			 SYS_TIMEOUT_MS(drv_data->tx.cfg.timeout));
	if (ret < 0) {
		return ret;
	}

	LOG_DBG("Queued TX %p", mem_block);

	/* Check if interrupt wanted to get next TX buffer before current buffer
	 * was queued. Do not move this check before queuing because doing so
	 * opens the possibility for a race condition between this function and
	 * data_handler() that is called in interrupt context.
	 */
	if (drv_data->state == I2S_STATE_RUNNING &&
	    drv_data->next_tx_buffer_needed) {
		nrfx_i2s_buffers_t next = { 0 };

		if (!get_next_tx_buffer(drv_data, &next)) {
			/* Log error because this is definitely unexpected.
			 * Do not return error because the caller is no longer
			 * responsible for releasing the buffer.
			 */
			LOG_ERR("Cannot reacquire queued buffer");
			return 0;
		}

		drv_data->next_tx_buffer_needed = false;

		LOG_DBG("Next TX %p", next.p_tx_buffer);

		if (!supply_next_buffers(drv_data, &next)) {
			return -EIO;
		}

	}

	return 0;
}

static int start_transfer(struct i2s_nrfx_drv_data *drv_data)
{
	nrfx_i2s_buffers_t initial_buffers = { 0 };
	int ret;

	if (drv_data->active_dir != I2S_DIR_RX && /* -> TX to be started */
	    !get_next_tx_buffer(drv_data, &initial_buffers)) {
		LOG_ERR("No TX buffer available");
		ret = -ENOMEM;
	} else if (drv_data->active_dir != I2S_DIR_TX && /* -> RX to be started */
		   !get_next_rx_buffer(drv_data, &initial_buffers)) {
		/* Failed to allocate next RX buffer */
		ret = -ENOMEM;
	} else {
		int err;

		/* It is necessary to set buffer size here only for I2S_DIR_RX,
		 * because only then the get_next_tx_buffer() call in the if
		 * condition above gets short-circuited.
		 */
		if (drv_data->active_dir == I2S_DIR_RX) {
			initial_buffers.buffer_size =
				drv_data->rx.cfg.block_size / sizeof(uint32_t);
		}

		drv_data->last_tx_buffer = initial_buffers.p_tx_buffer;

		err = nrfx_i2s_start(&drv_data->i2s, &initial_buffers, 0);
		if (err == 0) {
			return 0;
		}

		LOG_ERR("Failed to start I2S transfer: %d", err);
		ret = -EIO;
	}

	nrfx_i2s_uninit(&drv_data->i2s);
	if (drv_data->request_clock) {
		(void)onoff_release(drv_data->clk_mgr);
	}

	if (initial_buffers.p_tx_buffer) {
		free_tx_buffer(drv_data, initial_buffers.p_tx_buffer);
	}
	if (initial_buffers.p_rx_buffer) {
		free_rx_buffer(drv_data, initial_buffers.p_rx_buffer);
	}

	drv_data->state = I2S_STATE_ERROR;
	return ret;
}

static void clock_started_callback(struct onoff_manager *mgr,
				   struct onoff_client *cli,
				   uint32_t state,
				   int res)
{
	struct i2s_nrfx_drv_data *drv_data =
		CONTAINER_OF(cli, struct i2s_nrfx_drv_data, clk_cli);

	/* The driver state can be set back to READY at this point if the DROP
	 * command was triggered before the clock has started. Do not start
	 * the actual transfer in such case.
	 */
	if (drv_data->state == I2S_STATE_READY) {
		nrfx_i2s_uninit(&drv_data->i2s);
		(void)onoff_release(drv_data->clk_mgr);
	} else {
		(void)start_transfer(drv_data);
	}
}

static int trigger_start(const struct device *dev)
{
	struct i2s_nrfx_drv_data *drv_data = dev->data;
	const struct i2s_nrfx_drv_cfg *drv_cfg = dev->config;
	int err;
	int ret;
	const nrfx_i2s_config_t *nrfx_cfg = (drv_data->active_dir == I2S_DIR_TX)
					    ? &drv_data->tx.nrfx_cfg
					    : &drv_data->rx.nrfx_cfg;

	err = nrfx_i2s_init(&drv_data->i2s, nrfx_cfg, drv_cfg->data_handler);
	if (err != 0) {
		LOG_ERR("Failed to initialize I2S: %d", err);
		return -EIO;
	}

	drv_data->state = I2S_STATE_RUNNING;

#if NRF_I2S_HAS_CLKCONFIG
	nrf_i2s_clk_configure(drv_data->i2s.p_reg,
			      drv_cfg->clk_src == ACLK ? NRF_I2S_CLKSRC_ACLK
						       : NRF_I2S_CLKSRC_PCLK32M,
			      nrfx_cfg->prescalers.enable_bypass);
#endif

	/* If it is required to use certain HF clock, request it to be running
	 * first. If not, start the transfer directly.
	 */
	if (drv_data->request_clock) {
		sys_notify_init_callback(&drv_data->clk_cli.notify,
					 clock_started_callback);
		ret = onoff_request(drv_data->clk_mgr, &drv_data->clk_cli);
		if (ret < 0) {
			nrfx_i2s_uninit(&drv_data->i2s);
			drv_data->state = I2S_STATE_READY;

			LOG_ERR("Failed to request clock: %d", ret);
			return -EIO;
		}
	} else {
		ret = start_transfer(drv_data);
		if (ret < 0) {
			return ret;
		}
	}

	return 0;
}

static int i2s_nrfx_trigger(const struct device *dev,
			    enum i2s_dir dir, enum i2s_trigger_cmd cmd)
{
	struct i2s_nrfx_drv_data *drv_data = dev->data;
	bool configured = false;
	bool cmd_allowed;

	/* This driver does not use the I2S_STATE_NOT_READY value.
	 * Instead, if a given stream is not configured, the respective
	 * flag (tx_configured or rx_configured) is cleared.
	 */
	if (dir == I2S_DIR_BOTH) {
		configured = drv_data->tx_configured && drv_data->rx_configured;
	} else if (dir == I2S_DIR_TX) {
		configured = drv_data->tx_configured;
	} else if (dir == I2S_DIR_RX) {
		configured = drv_data->rx_configured;
	}

	if (!configured) {
		LOG_ERR("Device is not configured");
		return -EIO;
	}

	if (dir == I2S_DIR_BOTH &&
	    (memcmp(&drv_data->tx.nrfx_cfg,
		    &drv_data->rx.nrfx_cfg,
		    sizeof(drv_data->rx.nrfx_cfg)) != 0
	     ||
	     (drv_data->tx.cfg.block_size != drv_data->rx.cfg.block_size))) {
		LOG_ERR("TX and RX configurations are different");
		return -EIO;
	}

	switch (cmd) {
	case I2S_TRIGGER_START:
		cmd_allowed = (drv_data->state == I2S_STATE_READY);
		break;
	case I2S_TRIGGER_STOP:
	case I2S_TRIGGER_DRAIN:
		cmd_allowed = (drv_data->state == I2S_STATE_RUNNING);
		break;
	case I2S_TRIGGER_DROP:
		cmd_allowed = configured;
		break;
	case I2S_TRIGGER_PREPARE:
		cmd_allowed = (drv_data->state == I2S_STATE_ERROR);
		break;
	default:
		LOG_ERR("Invalid trigger: %d", cmd);
		return -EINVAL;
	}

	if (!cmd_allowed) {
		return -EIO;
	}

	/* For triggers applicable to the RUNNING state (i.e. STOP, DRAIN,
	 * and DROP), ensure that the command is applied to the streams
	 * that are currently active (this device cannot e.g. stop only TX
	 * without stopping RX).
	 */
	if (drv_data->state == I2S_STATE_RUNNING &&
	    drv_data->active_dir != dir) {
		LOG_ERR("Inappropriate trigger (%d/%d), active stream(s): %d",
			cmd, dir, drv_data->active_dir);
		return -EINVAL;
	}

	switch (cmd) {
	case I2S_TRIGGER_START:
		drv_data->stop = false;
		drv_data->discard_rx = false;
		drv_data->active_dir = dir;
		drv_data->next_tx_buffer_needed = false;
		return trigger_start(dev);

	case I2S_TRIGGER_STOP:
		drv_data->state = I2S_STATE_STOPPING;
		drv_data->stop = true;
		return 0;

	case I2S_TRIGGER_DRAIN:
		drv_data->state = I2S_STATE_STOPPING;
		/* If only RX is active, DRAIN is equivalent to STOP. */
		drv_data->stop = (drv_data->active_dir == I2S_DIR_RX);
		return 0;

	case I2S_TRIGGER_DROP:
		if (drv_data->state != I2S_STATE_READY) {
			drv_data->discard_rx = true;
			nrfx_i2s_stop(&drv_data->i2s);
		}
		purge_queue(dev, dir);
		drv_data->state = I2S_STATE_READY;
		return 0;

	case I2S_TRIGGER_PREPARE:
		purge_queue(dev, dir);
		drv_data->state = I2S_STATE_READY;
		return 0;

	default:
		LOG_ERR("Invalid trigger: %d", cmd);
		return -EINVAL;
	}
}

static void init_clock_manager(const struct device *dev)
{
	struct i2s_nrfx_drv_data *drv_data = dev->data;
	clock_control_subsys_t subsys;

#if NRF_CLOCK_HAS_HFCLKAUDIO
	const struct i2s_nrfx_drv_cfg *drv_cfg = dev->config;

	if (drv_cfg->clk_src == ACLK) {
		subsys = CLOCK_CONTROL_NRF_SUBSYS_HFAUDIO;
	} else
#endif
	{
		subsys = CLOCK_CONTROL_NRF_SUBSYS_HF;
	}

	drv_data->clk_mgr = z_nrf_clock_control_get_onoff(subsys);
	__ASSERT_NO_MSG(drv_data->clk_mgr != NULL);
}

static DEVICE_API(i2s, i2s_nrf_drv_api) = {
	.configure = i2s_nrfx_configure,
	.config_get = i2s_nrfx_config_get,
	.read = i2s_nrfx_read,
	.write = i2s_nrfx_write,
	.trigger = i2s_nrfx_trigger,
};

#define I2S_CLK_SRC(inst) DT_STRING_TOKEN(DT_DRV_INST(inst), clock_source)

#define I2S_NRFX_DEVICE(inst)                                                                      \
	static struct i2s_buf tx_msgs##inst[CONFIG_I2S_NRFX_TX_BLOCK_COUNT];                       \
	static struct i2s_buf rx_msgs##inst[CONFIG_I2S_NRFX_RX_BLOCK_COUNT];                       \
	static void data_handler##inst(nrfx_i2s_buffers_t const *p_released, uint32_t status)      \
	{                                                                                          \
		data_handler(DEVICE_DT_GET(DT_DRV_INST(inst)), p_released, status);                \
	}                                                                                          \
	PINCTRL_DT_DEFINE(DT_DRV_INST(inst));                                                      \
	static const struct i2s_nrfx_drv_cfg i2s_nrfx_cfg##inst = {                                \
		.data_handler = data_handler##inst,                                                \
		.nrfx_def_cfg = NRFX_I2S_DEFAULT_CONFIG(                                           \
			NRF_I2S_PIN_NOT_CONNECTED, NRF_I2S_PIN_NOT_CONNECTED,                      \
			NRF_I2S_PIN_NOT_CONNECTED, NRF_I2S_PIN_NOT_CONNECTED,                      \
			NRF_I2S_PIN_NOT_CONNECTED),                                                \
		.nrfx_def_cfg.skip_gpio_cfg = true,                                                \
		.nrfx_def_cfg.skip_psel_cfg = true,                                                \
		.pcfg = PINCTRL_DT_DEV_CONFIG_GET(DT_DRV_INST(inst)),                              \
		.clk_src = I2S_CLK_SRC(inst),                                                      \
	};                                                                                         \
	static struct i2s_nrfx_drv_data i2s_nrfx_data##inst = {                                    \
		.state = I2S_STATE_READY,                                                          \
		.i2s = NRFX_I2S_INSTANCE(DT_INST_REG_ADDR(inst)),                                  \
	};                                                                                         \
	static int i2s_nrfx_init##inst(const struct device *dev)                                   \
	{                                                                                          \
		IRQ_CONNECT(DT_INST_IRQN(inst), DT_INST_IRQ(inst, priority), nrfx_i2s_irq_handler, \
			    &i2s_nrfx_data##inst.i2s, 0);                                          \
		const struct i2s_nrfx_drv_cfg *drv_cfg = dev->config;                              \
		int err = pinctrl_apply_state(drv_cfg->pcfg, PINCTRL_STATE_DEFAULT);               \
		if (err < 0) {                                                                     \
			return err;                                                                \
		}                                                                                  \
		k_msgq_init(&i2s_nrfx_data##inst.tx_queue, (char *)tx_msgs##inst,                  \
			    sizeof(struct i2s_buf), ARRAY_SIZE(tx_msgs##inst));                    \
		k_msgq_init(&i2s_nrfx_data##inst.rx_queue, (char *)rx_msgs##inst,                  \
			    sizeof(struct i2s_buf), ARRAY_SIZE(rx_msgs##inst));                    \
		init_clock_manager(dev);                                                           \
		return 0;                                                                          \
	}                                                                                          \
	BUILD_ASSERT(I2S_CLK_SRC(inst) != ACLK ||                                                  \
			     (NRF_I2S_HAS_CLKCONFIG && NRF_CLOCK_HAS_HFCLKAUDIO),                  \
		     "Clock source ACLK is not available.");                                       \
	BUILD_ASSERT(I2S_CLK_SRC(inst) != ACLK ||                                                  \
			     DT_NODE_HAS_PROP(DT_NODELABEL(clock), hfclkaudio_frequency),          \
		     "Clock source ACLK requires the hfclkaudio-frequency "                        \
		     "property to be defined in the nordic,nrf-clock node.");                      \
	DEVICE_DT_INST_DEFINE(inst, i2s_nrfx_init##inst, NULL, &i2s_nrfx_data##inst,               \
			      &i2s_nrfx_cfg##inst, POST_KERNEL, CONFIG_I2S_INIT_PRIORITY,          \
			      &i2s_nrf_drv_api);

DT_INST_FOREACH_STATUS_OKAY(I2S_NRFX_DEVICE)
